Access values in dict in tf.dataset.Dataset.map with tf.striing